package org.easy.milvus;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import com.alibaba.fastjson2.annotation.JSONField;
import org.easy.core.util.HttpUtil;
import org.easy.core.util.Lists;
import lombok.AllArgsConstructor;
import lombok.Data;

import java.util.ArrayList;
import java.util.List;


/**
 * @author CYY
 * @date 2023年03月11日 下午9:00
 * @description
 */

public class EmbeddingClient {

    public EmbeddingClient(String url) {
        this.url = url;
    }

    /**
     * http://localhost:8001/get_text_embedding
     */
    private String url;

    /**
     * 请求获取Embeddings，请求出错返回null
     *
     * @param msg 需要Embeddings的信息
     * @return 为null则请求失败，反之放回正确结果
     */
    public List<List<Float>> doEmbedding(String msg) {

        ArrayList<String> sourceSentence = Lists.newArrayList(msg);

        return doEmbedding(sourceSentence);
    }

    public List<List<Float>> doEmbedding(List<String> sourceSentence) {
        EmbeddingInput embeddingInput = new EmbeddingInput(sourceSentence);
        String jsonArrayStr = HttpUtil.postJson(url, null, embeddingInput);
        return JSON.parseObject(jsonArrayStr, new TypeReference<List<List<Float>>>() {
        });
    }


    @Data
    @AllArgsConstructor
    public class EmbeddingInput {
        @JSONField(name = "source_sentence")
        private List<String> sourceSentence;
    }

}
